extern "C" {
#include <libavcodec/avcodec.h>
#include <libavdevice/avdevice.h>
#include <libavformat/avformat.h>
#include <libavutil/imgutils.h>
#include <libavutil/samplefmt.h>
#include <libavutil/timestamp.h>
}
#include "window.h"
#include <cassert>
#include <iostream>
using namespace std;

#ifdef av_err2str
#undef av_err2str
#include <string>
av_always_inline std::string av_err2string(int errnum) {
    char str[AV_ERROR_MAX_STRING_SIZE];
    return av_make_error_string(str, AV_ERROR_MAX_STRING_SIZE, errnum);
}
#define av_err2str(err) av_err2string(err).c_str()
#endif // av_err2str

static AVFormatContext *fmt_ctx = NULL;
static AVCodecContext *video_dec_ctx = NULL, *audio_dec_ctx;
static int width, height;
static enum AVPixelFormat pix_fmt;
static AVStream *video_stream = NULL, *audio_stream = NULL;
static const char *src_filename = NULL;

static FILE* video_dst_file = NULL;
static uint8_t *video_dst_data[4] = {NULL};
static int video_dst_linesize[4];
static int video_dst_bufsize;

static int video_stream_idx = -1, audio_stream_idx = -1;
static AVFrame *frame = NULL;
static AVPacket *pkt = NULL;
static int video_frame_count = 0;
static int audio_frame_count = 0;

static int output_video_frame(AVFrame *frame) {
    if (frame->width != width || frame->height != height ||
        frame->format != pix_fmt) {
        /* To handle this change, one could call av_image_alloc again and
         * decode the following frames into another rawvideo file. */
        fprintf(stderr,
                "Error: Width, height and pixel format have to be "
                "constant in a rawvideo file, but the width, height or "
                "pixel format of the input video changed:\n"
                "old: width = %d, height = %d, format = %s\n"
                "new: width = %d, height = %d, format = %s\n",
                width, height, av_get_pix_fmt_name(pix_fmt), frame->width,
                frame->height,
                av_get_pix_fmt_name((AVPixelFormat)frame->format));
        //return -1;
    }

    cout << av_get_pix_fmt_name((AVPixelFormat)frame->format) << endl;
    printf("video_frame n:%d coded_n:%d\n", video_frame_count++,
           frame->coded_picture_number);

    /* copy decoded frame to destination buffer:
     * this is required since rawvideo expects non aligned data */
    av_image_copy(video_dst_data, video_dst_linesize,
                  (const uint8_t **)(frame->data), frame->linesize, pix_fmt,
                  width, height);

     //fwrite(video_dst_data[0], 1, video_dst_bufsize, video_dst_file);

    window::getInstance().fillVideo(video_dst_data, frame);
    return 0;
}

static int output_audio_frame(AVFrame *frame) {
    size_t unpadded_linesize =
        frame->nb_samples *
        av_get_bytes_per_sample((AVSampleFormat)frame->format);
    // printf("audio_frame n:%d nb_samples:%d pts:%s\n",
    // audio_frame_count++, frame->nb_samples,
    // av_ts2timestr(frame->pts, &audio_dec_ctx->time_base));

    /* Write the raw audio data samples of the first plane. This works
     * fine for packed formats (e.g. AV_SAMPLE_FMT_S16). However,
     * most audio decoders output planar audio, which uses a separate
     * plane of audio samples for each channel (e.g. AV_SAMPLE_FMT_S16P).
     * In other words, this code will write only the first audio channel
     * in these cases.
     * You should use libswresample or libavfilter to convert the frame
     * to packed data. */
    //fwrite(frame->extended_data[0], 1, unpadded_linesize, audio_dst_file);

    return 0;
}

static int decode_packet(AVCodecContext *dec, const AVPacket *pkt) {
    int ret = 0;

    // submit the packet to the decoder
    ret = avcodec_send_packet(dec, pkt);
    if (ret < 0) {
        fprintf(stderr, "Error submitting a packet for decoding (%s)\n",
                av_err2str(ret));
        return ret;
    }

    // get all the available frames from the decoder
    while (ret >= 0) {
        window::getInstance().control();
        ret = avcodec_receive_frame(dec, frame);
        if (ret < 0) {
            // those two return values are special and mean there is no output
            // frame available, but there were no errors during decoding
            if (ret == AVERROR_EOF || ret == AVERROR(EAGAIN))
                return 0;

            fprintf(stderr, "Error during decoding (%s)\n", av_err2str(ret));
            return ret;
        }

        // write the frame data to output file
        if (dec->codec->type == AVMEDIA_TYPE_VIDEO)
            ret = output_video_frame(frame);
        // else
        // ret = output_audio_frame(frame);

        av_frame_unref(frame);
        if (ret < 0)
            return ret;
    }

    return 0;
}

static int open_codec_context(int *stream_idx, AVCodecContext **dec_ctx,
                              AVFormatContext *fmt_ctx, enum AVMediaType type) {
    int ret, stream_index;
    AVStream *st;
    const AVCodec *dec = NULL;

    ret = av_find_best_stream(fmt_ctx, type, -1, -1, NULL, 0);
    if (ret < 0) {
        fprintf(stderr, "Could not find %s stream in input file '%s'\n",
                av_get_media_type_string(type), src_filename);
        return ret;
    } else {
        stream_index = ret;
        st = fmt_ctx->streams[stream_index];

        /* find decoder for the stream */
        dec = avcodec_find_decoder(st->codecpar->codec_id);
        if (!dec) {
            fprintf(stderr, "Failed to find %s codec\n",
                    av_get_media_type_string(type));
            return AVERROR(EINVAL);
        }

        /* Allocate a codec context for the decoder */
        *dec_ctx = avcodec_alloc_context3(dec);
        if (!*dec_ctx) {
            fprintf(stderr, "Failed to allocate the %s codec context\n",
                    av_get_media_type_string(type));
            return AVERROR(ENOMEM);
        }

        /* Copy codec parameters from input stream to output codec context */
        if ((ret = avcodec_parameters_to_context(*dec_ctx, st->codecpar)) < 0) {
            fprintf(stderr,
                    "Failed to copy %s codec parameters to decoder context\n",
                    av_get_media_type_string(type));
            return ret;
        }

        /* Init the decoders */
        if ((ret = avcodec_open2(*dec_ctx, dec, NULL)) < 0) {
            fprintf(stderr, "Failed to open %s codec\n",
                    av_get_media_type_string(type));
            return ret;
        }
        *stream_idx = stream_index;
    }

    return 0;
}

static int get_format_from_sample_fmt(const char **fmt,
                                      enum AVSampleFormat sample_fmt) {
    int i;
    struct sample_fmt_entry {
        enum AVSampleFormat sample_fmt;
        const char *fmt_be, *fmt_le;
    } sample_fmt_entries[] = {
        {AV_SAMPLE_FMT_U8, "u8", "u8"},
        {AV_SAMPLE_FMT_S16, "s16be", "s16le"},
        {AV_SAMPLE_FMT_S32, "s32be", "s32le"},
        {AV_SAMPLE_FMT_FLT, "f32be", "f32le"},
        {AV_SAMPLE_FMT_DBL, "f64be", "f64le"},
    };
    *fmt = NULL;

    for (i = 0; i < FF_ARRAY_ELEMS(sample_fmt_entries); i++) {
        struct sample_fmt_entry *entry = &sample_fmt_entries[i];
        if (sample_fmt == entry->sample_fmt) {
            *fmt = AV_NE(entry->fmt_be, entry->fmt_le);
            return 0;
        }
    }

    fprintf(stderr, "sample format %s is not supported as output format\n",
            av_get_sample_fmt_name(sample_fmt));
    return -1;
}

int test_decode_video(int argc, char **argv) {
    int ret = 0;
    src_filename = "flv.flv";
    video_dst_file = fopen("out.raw", "wb");

    /* open input file, and allocate format context */
    if (avformat_open_input(&fmt_ctx, src_filename, NULL, NULL) < 0) {
        fprintf(stderr, "Could not open source file %s\n", src_filename);
        exit(1);
    }

    /* retrieve stream information */
    if (avformat_find_stream_info(fmt_ctx, NULL) < 0) {
        fprintf(stderr, "Could not find stream information\n");
        exit(1);
    }

    if (open_codec_context(&video_stream_idx, &video_dec_ctx, fmt_ctx,
                           AVMEDIA_TYPE_VIDEO) >= 0) {
        video_stream = fmt_ctx->streams[video_stream_idx];
        width = video_dec_ctx->width;
        height = video_dec_ctx->height;
        pix_fmt = video_dec_ctx->pix_fmt;
        cout << "Width: " << width << " Height: " << height << endl;
        ret = av_image_alloc(video_dst_data, video_dst_linesize,
                width, height, pix_fmt, 1);
        if (ret < 0) {
            fprintf(stderr, "Could not allocate raw video buffer\n");
            goto end;
        }
        video_dst_bufsize = ret;
        window::getInstance().setWH(width, height, pix_fmt);
    }

    if (open_codec_context(&audio_stream_idx, &audio_dec_ctx, fmt_ctx, AVMEDIA_TYPE_AUDIO) >= 0) {
        audio_stream = fmt_ctx->streams[audio_stream_idx];
    }

    if (!audio_stream && !video_stream) {
        fprintf(
            stderr,
            "Could not find audio or video stream in the input, aborting\n");
        ret = 1;
        goto end;
    }

    frame = av_frame_alloc();
    pkt = av_packet_alloc();
    if (video_stream)
        printf("Demuxing video from file '%s' into 'dummy'\n", src_filename);
    if (audio_stream)
        printf("Demuxing audio from file '%s' into 'dummy'\n", src_filename);

    /* read frames from the file */
    while (av_read_frame(fmt_ctx, pkt) >= 0) {
        // check if the packet belongs to a stream we are interested in,
        // otherwise skip it
        if (pkt->stream_index == video_stream_idx)
            ret = decode_packet(video_dec_ctx, pkt);
        else if (pkt->stream_index == audio_stream_idx)
            ret = decode_packet(audio_dec_ctx, pkt);
        av_packet_unref(pkt);
        if (ret < 0)
            break;
    }

    /* flush the decoders */
    if (video_dec_ctx)
        decode_packet(video_dec_ctx, NULL);
    if (audio_dec_ctx)
        decode_packet(audio_dec_ctx, NULL);

    printf("Demuxing succeeded.\n");

    if (video_stream) {
        printf("Play the output video file with the command:\n"
               "ffplay -f rawvideo -pix_fmt %s -video_size %dx%d dummy\n",
               av_get_pix_fmt_name(pix_fmt), width, height);
    }

    if (audio_stream) {
        enum AVSampleFormat sfmt = audio_dec_ctx->sample_fmt;
        int n_channels = audio_dec_ctx->ch_layout.nb_channels;
        const char *fmt;

        if (av_sample_fmt_is_planar(sfmt)) {
            const char *packed = av_get_sample_fmt_name(sfmt);
            printf("Warning: the sample format the decoder produced is planar "
                   "(%s). This example will output the first channel only.\n",
                   packed ? packed : "?");
            sfmt = av_get_packed_sample_fmt(sfmt);
            n_channels = 1;
        }

        if ((ret = get_format_from_sample_fmt(&fmt, sfmt)) < 0)
            goto end;

        printf("Play the output audio file with the command:\n"
               "ffplay -f %s -ac %d -ar %d dummy\n",
               fmt, n_channels, audio_dec_ctx->sample_rate);
    }

end:
    avcodec_free_context(&video_dec_ctx);
    avcodec_free_context(&audio_dec_ctx);
    avformat_close_input(&fmt_ctx);
    av_packet_free(&pkt);
    av_frame_free(&frame);
    av_free(video_dst_data[0]);
    if (video_dst_file) {
        fclose(video_dst_file);
    }

    return ret < 0;
}

int gdigrab_stream() {
    avdevice_register_all();
    const AVInputFormat *ifmt = av_find_input_format("gdigrab");
    if (!ifmt) {
        cout << "find input format error" << endl;
        return -1;
    }
    AVDictionary *options = NULL;
    av_dict_set(&options, "framerate", "30", 0);
    av_dict_set(&options, "offset_x", "0", 0);
    av_dict_set(&options, "offset_y", "0", 0);
    //av_dict_set(&options, "video_size", "1920x1080", 0);
    AVFormatContext *ifmt_ctx = NULL;
    int r = 0;
    if ((r = avformat_open_input(&ifmt_ctx, "desktop", ifmt, &options)) < 0) {
        cerr << av_err2str(r) << endl;
        std::cerr << "Failed to open input device" << std::endl;
        return -1;
    }

    // Find the video stream information
    if (avformat_find_stream_info(ifmt_ctx, NULL) < 0) {
        std::cerr << "Failed to retrieve input stream information" << std::endl;
        return -1;
    }

    // Find the video stream
    int video_stream_index = -1;
    cout << "ifmt_ctx.nb_streams: " << ifmt_ctx->nb_streams << endl;
    for (int i = 0; i < ifmt_ctx->nb_streams; i++) {
        if (ifmt_ctx->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO) {
            video_stream_index = i;
            break;
        }
    }
    if (video_stream_index == -1) {
        std::cerr << "Failed to find video stream" << std::endl;
        return -1;
    }
    cout << "Video Stream Index: " << video_stream_index << endl;
    
    int ret = 0;
    if (open_codec_context(&video_stream_idx, &video_dec_ctx, ifmt_ctx, AVMEDIA_TYPE_VIDEO) >= 0) {
        video_stream = ifmt_ctx->streams[video_stream_index];
        width = video_dec_ctx->width;
        height = video_dec_ctx->height;
        pix_fmt = video_dec_ctx->pix_fmt;
        cout << "Width: " << width << " Height: " << height << "PixFormat: "
                << av_get_pix_fmt_name(pix_fmt) << endl;
        ret = av_image_alloc(video_dst_data, video_dst_linesize,
                width, height, pix_fmt, 1);
        if (ret < 0) {
            fprintf(stderr, "Could not allocate raw video buffer\n");
            return -1;
        }
        video_dst_bufsize = ret;
        window::getInstance().setWH(width, height, pix_fmt);
    } else {
        return -1;
    }

    if (video_stream) {
        // 获取gdigrab的数据
        // Read frames from the window capture device and write them to the
        // output file
        pkt = av_packet_alloc();
        frame = av_frame_alloc();
        int ret = 0;
        while (av_read_frame(ifmt_ctx, pkt) >= 0) {
            if (pkt->stream_index == video_stream_index)
                ret = decode_packet(video_dec_ctx, pkt);
            av_packet_unref(pkt);
            cout << ret << endl;
            if (ret < 0) {
                break;
            }
        }
    }

    if (video_dec_ctx)
        decode_packet(video_dec_ctx, NULL);
    avcodec_free_context(&video_dec_ctx);
    avformat_close_input(&ifmt_ctx);
    av_frame_free(&frame);
    av_packet_free(&pkt);
    av_free(video_dst_data[0]);
}

// 桌面捕捉 gdigrab
void test_avdevice() {
    avdevice_register_all();
    const AVInputFormat *d = NULL;
    do {
        d = av_input_video_device_next(d);
        if (!d) {
            break;
        }
        // windows supoort video input devices
        // dshow
        // gdigrab
        // lavfi
        // vfwcap
        if (string(d->name) == "dshow") {
            cout << d->name << endl;
            AVDeviceInfoList *infos[10] = {0};
            avdevice_list_input_sources(d, NULL, NULL, infos);
            int i = 0;
            while (infos[i]) {
                cout << infos[i]->nb_devices << endl;
                int k = 0;
                AVDeviceInfo *deviceinfo = infos[i]->devices[k];
                while (deviceinfo) {
                    cout << k << ". ";
                    cout << deviceinfo->device_description << endl;
                    cout << deviceinfo->device_name << endl;
                    cout << endl;
                    deviceinfo = infos[i]->devices[k++];
                }
                i++;
            }
        }
    } while (d);
}

int main(int argc, char *argv[]) {
     //test_decode_video(argc, argv);
    // test_avdevice();
    gdigrab_stream();
    return 0;
}
